import os
from torch.utils.data import Dataset
from PIL import Image
from torchvision import transforms

data_transform = {
    "train": transforms.Compose([transforms.ToTensor(),
                                 transforms.Normalize((0.5,), (0.5,))]),
    "test": transforms.Compose([transforms.ToTensor()])
}

class Data_Loader(Dataset):
    def __init__(self, root, transforms_train=data_transform['train'], transforms_test=data_transform['test']):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]
        self.transforms_train = transforms_train
        self.transforms_test = transforms_test

    def __getitem__(self, index):
        image_path = self.imgs[index]
        # 生成标签路径
        # 获取文件名
        base_name = os.path.basename(image_path)
        # 去掉扩展名
        base_name = os.path.splitext(base_name)[0]
        # 移除 'training_gray_' 字段
        base_name = base_name.replace('_training_gray', '')
        label_name = f"{base_name}_manual1.png"
        label_path = os.path.join(os.path.dirname(image_path).replace('images', 'label'), label_name)

        image = Image.open(image_path)
        label = Image.open(label_path)

        image = self.transforms_train(image)
        label = self.transforms_test(label)
        label[label > 0] = 1

        return image, label

    def __len__(self):
        return len(self.imgs)